#!/bin/bash
set -e

echo "Applying NPU patches to HunyuanWorld-Mirror code..."

TARGET_DIR="/app/HunyuanWorld-Mirror"

if [ ! -d "$TARGET_DIR" ]; then
    echo "Error: Target directory $TARGET_DIR not found."
    exit 1
fi

cd "$TARGET_DIR"

# 1. 替换 infer.py 中的 CUDA 调用
if [ -f "infer.py" ]; then
    echo "Patching infer.py..."
    # 添加 import torch_npu (如果还没有)
    grep -q "import torch_npu" infer.py || sed -i '1s/^/import torch_npu\n/' infer.py
    
    # 替换 device 设置
    sed -i 's/device = torch.device("cuda" if torch.cuda.is_available() else "cpu")/device = torch.device("npu" if torch.npu.is_available() else "cpu")/' infer.py
    
    # 禁用混合精度 (NPU 上 scatter_add 不支持 bfloat16)
    sed -i 's/use_amp = torch.cuda.is_available() and torch.cuda.is_bf16_supported()/use_amp = False  # Disabled for NPU compatibility/' infer.py
    
    # 替换 autocast
    sed -i "s/torch.amp.autocast('cuda'/torch.amp.autocast('npu'/" infer.py
fi

# 2. 替换 app.py 中的 CUDA 调用
if [ -f "app.py" ]; then
    echo "Patching app.py..."
    # 添加 import torch_npu (如果还没有)
    grep -q "import torch_npu" app.py || sed -i '1s/^/import torch_npu\n/' app.py
    
    # 移除 PYTORCH_CUDA_ALLOC_CONF 环境变量设置
    sed -i '/os.environ\["PYTORCH_CUDA_ALLOC_CONF"\]/d' app.py
    
    # 替换 device 设置
    sed -i 's/device = "cuda" if torch.cuda.is_available() else "cpu"/device = "npu" if torch.npu.is_available() else "cpu"/' app.py
    
    # 禁用混合精度 (NPU 上某些算子不支持 bfloat16)
    sed -i 's/use_amp = torch.cuda.is_available() and torch.cuda.is_bf16_supported()/use_amp = False  # Disabled for NPU compatibility/' app.py
    
    # 替换 autocast
    sed -i "s/torch.amp.autocast('cuda'/torch.amp.autocast('npu'/" app.py
    
    # 替换 torch.cuda.empty_cache()
    sed -i 's/torch.cuda.empty_cache()/torch.npu.empty_cache()/g' app.py
fi

# 3. 修复 src/models/models 目录下的文件
echo "Patching src/models/models/*.py..."

# 修复 worldmirror.py
if [ -f "src/models/models/worldmirror.py" ]; then
    grep -q "import torch_npu" src/models/models/worldmirror.py || sed -i '1s/^/try:\n    import torch_npu\nexcept ImportError:\n    pass\n/' src/models/models/worldmirror.py
fi

# 修复 rasterization.py - 关键文件，包含 scatter_add 操作
if [ -f "src/models/models/rasterization.py" ]; then
    echo "Patching rasterization.py for NPU compatibility..."
    grep -q "import torch_npu" src/models/models/rasterization.py || sed -i '1s/^/try:\n    import torch_npu\nexcept ImportError:\n    pass\n/' src/models/models/rasterization.py
    
    # 将 scatter_add_ 操作前的张量转换为 float32 (如果是 bfloat16 会出错)
    # 这需要在运行时处理，脚本无法完美修复，但可以添加注释提醒
fi

# 4. 修复 src/utils 目录下的文件
echo "Patching src/utils/*.py..."

if [ -f "src/utils/visual_util.py" ]; then
    grep -q "import torch_npu" src/utils/visual_util.py || sed -i '1s/^/try:\n    import torch_npu\nexcept ImportError:\n    pass\n/' src/utils/visual_util.py
fi

# 5. 全局替换所有 Python 文件中的 cuda 调用
echo "Global CUDA -> NPU replacements..."

# 在所有 Python 文件中替换 .cuda() 调用为 .npu() (谨慎操作)
find . -name "*.py" -type f -exec grep -l "\.cuda()" {} \; | while read file; do
    echo "  Patching $file..."
    sed -i 's/\.cuda()/.npu()/g' "$file"
done

# 替换 torch.cuda.is_available() 为 torch.npu.is_available()
find . -name "*.py" -type f -exec grep -l "torch.cuda.is_available()" {} \; | while read file; do
    echo "  Patching cuda.is_available in $file..."
    sed -i 's/torch.cuda.is_available()/torch.npu.is_available()/g' "$file"
done

# 替换 torch.cuda.empty_cache() 为 torch.npu.empty_cache()
find . -name "*.py" -type f -exec grep -l "torch.cuda.empty_cache()" {} \; | while read file; do
    echo "  Patching cuda.empty_cache in $file..."
    sed -i 's/torch.cuda.empty_cache()/torch.npu.empty_cache()/g' "$file"
done

echo ""
echo "=========================================="
echo "NPU patching completed!"
echo ""
echo "IMPORTANT: Some operations like scatter_add may still fail"
echo "with bfloat16 dtype. If you see dtype errors, the API has"
echo "been configured to use float32 instead of bfloat16."
echo "=========================================="
